sqlitevec: use DELETE by key instead of IN for virtual table deletes - #53

// Benchmarks for https://github.com/CommunityToolkit/AI/issues/52
// Compares SqliteVec batch delete/upsert cost on the vec0 virtual table.
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Configs;
using BenchmarkDotNet.Running;
using CommunityToolkit.VectorData.SqliteVec;
using Microsoft.Extensions.VectorData;
namespace SqliteVecBench;
public sealed class Record
{
    [VectorStoreKey(StorageName = "chunk_id")]
    public string ChunkId { get; set; }
    [VectorStoreData]
    public string Text { get; set; }
    [VectorStoreVector(Dimensions: 256, DistanceFunction = DistanceFunction.CosineDistance)]
    public ReadOnlyMemory<float> Embedding { get; set; }
}
[MemoryDiagnoser(displayGenColumns: false)]
public class SqliteVecDeleteBenchmarks
{
    private const int BatchSize = 200;
    private string _dbPath;
    private SqliteCollection<string, Record> _collection;
    private List<string> _missingKeys;
    private List<string> _existingKeys;
    private List<Record> _existingRecords;
    [Params(10_000, 100_000)]
    public int RowCount { get; set; }
    private static ReadOnlyMemory<float> CreateVector(Random random)
    {
        float[] values = new float[256];
        for (int i = 0; i < values.Length; i++)
        {
            values[i] = (float)random.NextDouble();
        }
        return new ReadOnlyMemory<float>(values);
    }
    [GlobalSetup]
    public async Task SetupAsync()
    {
        _dbPath = Path.Combine(Path.GetTempPath(), $"sqlitevec-bench-{RowCount}-{Guid.NewGuid():N}.db");
        _collection = new SqliteCollection<string, Record>($"Data Source={_dbPath}", "vec_chunks");
        await _collection.EnsureCollectionExistsAsync();
        Random random = new Random(42);
        List<Record> batch = new List<Record>(1000);
        for (int i = 0; i < RowCount; i++)
        {
            batch.Add(new Record { ChunkId = $"key-{i}", Text = "text", Embedding = CreateVector(random) });
            if (batch.Count == 1000)
            {
                await _collection.UpsertAsync(batch);
                batch.Clear();
            }
        }
        if (batch.Count > 0)
        {
            await _collection.UpsertAsync(batch);
        }
        _missingKeys = Enumerable.Range(0, BatchSize).Select(i => $"missing-{i}").ToList();
        // Keys spread across the table, from the "middle" of the key space.
        _existingKeys = Enumerable.Range(0, BatchSize).Select(i => $"key-{i * (RowCount / BatchSize)}").ToList();
        _existingRecords = _existingKeys
            .Select(k => new Record { ChunkId = k, Text = "text", Embedding = CreateVector(random) })
            .ToList();
    }
    [GlobalCleanup]
    public void Cleanup()
    {
        _collection?.Dispose();
        Microsoft.Data.Sqlite.SqliteConnection.ClearAllPools();
        if (_dbPath is not null && File.Exists(_dbPath))
        {
            File.Delete(_dbPath);
        }
    }
    // Delete a batch of keys that are not in the table: no rows are removed, so the
    // benchmark is idempotent and measures the lookup cost only.
    [Benchmark]
    public Task DeleteBatch_MissingKeys() => _collection.DeleteAsync(_missingKeys);
    // Single key delete for a key that is not present.
    [Benchmark]
    public Task DeleteSingle_MissingKey() => _collection.DeleteAsync("missing-0");
    // Upsert of records that already exist: internally deletes the vector rows of the
    // batch and re-inserts them, so the table size stays constant.
    [Benchmark]
    public Task UpsertBatch_ExistingRecords() => _collection.UpsertAsync(_existingRecords);
}

public static class Program
{
    public static void Main(string[] args)
        => BenchmarkSwitcher.FromAssembly(typeof(Program).Assembly).Run(args);
}

// Benchmarks for https://github.com/CommunityToolkit/AI/issues/52
// Compares SqliteVec batch delete/upsert cost on the vec0 virtual table.
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Configs;
using BenchmarkDotNet.Running;
using CommunityToolkit.VectorData.SqliteVec;
using Microsoft.Extensions.VectorData;
namespace SqliteVecBench;
public sealed class Record
{
    [VectorStoreKey(StorageName = "chunk_id")]
    public string ChunkId { get; set; }
    [VectorStoreData]
    public string Text { get; set; }
    [VectorStoreVector(Dimensions: 256, DistanceFunction = DistanceFunction.CosineDistance)]
    public ReadOnlyMemory<float> Embedding { get; set; }
}
[MemoryDiagnoser(displayGenColumns: false)]
public class SqliteVecDeleteBenchmarks
{
    private const int BatchSize = 200;
    private string _dbPath;
    private SqliteCollection<string, Record> _collection;
    private List<string> _missingKeys;
    private List<string> _existingKeys;
    private List<Record> _existingRecords;
    [Params(10_000, 100_000)]
    public int RowCount { get; set; }
    private static ReadOnlyMemory<float> CreateVector(Random random)
    {
        float[] values = new float[256];
        for (int i = 0; i < values.Length; i++)
        {
            values[i] = (float)random.NextDouble();
        }
        return new ReadOnlyMemory<float>(values);
    }
    [GlobalSetup]
    public async Task SetupAsync()
    {
        _dbPath = Path.Combine(Path.GetTempPath(), $"sqlitevec-bench-{RowCount}-{Guid.NewGuid():N}.db");
        _collection = new SqliteCollection<string, Record>($"Data Source={_dbPath}", "vec_chunks");
        await _collection.EnsureCollectionExistsAsync();
        Random random = new Random(42);
        List<Record> batch = new List<Record>(1000);
        for (int i = 0; i < RowCount; i++)
        {
            batch.Add(new Record { ChunkId = $"key-{i}", Text = "text", Embedding = CreateVector(random) });
            if (batch.Count == 1000)
            {
                await _collection.UpsertAsync(batch);
                batch.Clear();
            }
        }
        if (batch.Count > 0)
        {
            await _collection.UpsertAsync(batch);
        }
        _missingKeys = Enumerable.Range(0, BatchSize).Select(i => $"missing-{i}").ToList();
        // Keys spread across the table, from the "middle" of the key space.
        _existingKeys = Enumerable.Range(0, BatchSize).Select(i => $"key-{i * (RowCount / BatchSize)}").ToList();
        _existingRecords = _existingKeys
            .Select(k => new Record { ChunkId = k, Text = "text", Embedding = CreateVector(random) })
            .ToList();
    }
    [GlobalCleanup]
    public void Cleanup()
    {
        _collection?.Dispose();
        Microsoft.Data.Sqlite.SqliteConnection.ClearAllPools();
        if (_dbPath is not null && File.Exists(_dbPath))
        {
            File.Delete(_dbPath);
        }
    }
    // Delete a batch of keys that are not in the table: no rows are removed, so the
    // benchmark is idempotent and measures the lookup cost only.
    [Benchmark]
    public Task DeleteBatch_MissingKeys() => _collection.DeleteAsync(_missingKeys);
    // Single key delete for a key that is not present.
    [Benchmark]
    public Task DeleteSingle_MissingKey() => _collection.DeleteAsync("missing-0");
    // Upsert of records that already exist: internally deletes the vector rows of the
    // batch and re-inserts them, so the table size stays constant.
    [Benchmark]
    public Task UpsertBatch_ExistingRecords() => _collection.UpsertAsync(_existingRecords);
}

public static class Program
{
    public static void Main(string[] args)
        => BenchmarkSwitcher.FromAssembly(typeof(Program).Assembly).Run(args);
}